import { AgentInfo } from '../Agent';
import { GrpcExtensions } from '../Communicator/GrpcExtensions';
import { BrainParameters } from '../Policies/BrainParameters';
import { ISensor } from '../Sensors/ISensor';
import { ObservationWriter } from '../Sensors/ObservationWriter';
import { DemonstrationMetaData } from './DemonstrationMetaData';

import pb from '../Grpc/CommunicatorObjects/pb.js';
import * as protobuf from 'protobufjs';

export class DemonstrationWriter {
	/**
	 * @en
	 * Number of bytes reserved for the {@link DemonstrationMetaData} at the start of the demo file.
	 *
	 * @zh
	 * 在演示文件开头为 {@link DemonstrationMetaData} 保留的字节数。
	 */
	private static readonly MetaDataBytes: number = 32;

	private _metaData: DemonstrationMetaData = null!;
	private _writerName: string | null = null;
	private _cumulativeReward: number = 0;
	private _observationWriter: ObservationWriter = new ObservationWriter();

	constructor(writerName: string) {
		this._writerName = writerName;
		this.openWriter();
	}

	get numSteps(): number {
		return this._metaData.numberSteps;
	}

	public initialize(demonstrationName: string, brainParameters: BrainParameters, brainName: string): void {
		if (this._writerName == null) {
			return;
		}

		this._metaData = new DemonstrationMetaData();
		this._metaData.demonstrationName = demonstrationName;
		const metaProto = GrpcExtensions.DemonstrationMetaDataToProto(this._metaData);
		const writer = protobuf.Writer.create();
		pb.communicator_objects.DemonstrationMetaProto.encode(metaProto, writer);
		const binaryArray = writer.finish();
		this.write(binaryArray);
		this.writeBrainParameters(brainName, brainParameters);
	}

	private openWriter() {
		Editor.Message.send('mlagents', 'openDemonstrationWriter', this._writerName);
	}

	private write(binaryArray: Uint8Array) {
		const buffer = Buffer.from(binaryArray);
		Editor.Message.send('mlagents', 'writeDemonstrationBuffer', this._writerName, buffer);
	}

	private closeWriter() {
		Editor.Message.send('mlagents', 'closeDemonstrationWriter', this._writerName);
	}

	private seek(offset: number) {
		Editor.Message.send('mlagents', 'seekDemonstrationWriter', this._writerName, offset);
	}

	private writeMetadata(): void {
		if (this._writerName == null) {
			return;
		}
		const writer = protobuf.Writer.create();
		const metaProto = GrpcExtensions.DemonstrationMetaDataToProto(this._metaData);
		pb.communicator_objects.DemonstrationMetaProto.encode(metaProto, writer);
		const metaProtoBytes = writer.finish();
		this.write(metaProtoBytes);
		this.seek(0);

		const binaryArray = writer.finish();
		this.write(binaryArray);
	}

	private writeBrainParameters(brainName: string, brainParameters: BrainParameters): void {
		if (this._writerName == null) {
			return;
		}

		this.seek(DemonstrationWriter.MetaDataBytes + 1);
		const brainProto = GrpcExtensions.brainParametersToProto(brainParameters, brainName, false);
		const binaryArray = pb.communicator_objects.BrainParametersProto.encode(brainProto).toArray();
		this.write(binaryArray);
	}

	public record(info: AgentInfo, sensors: ISensor[]): void {
		if (this._writerName == null) {
			return;
		}

		this._metaData.numberSteps++;
		this._cumulativeReward += info.reward;
		if (info.done) {
			this.endEpisode();
		}

		const agentProto = GrpcExtensions.toInfoActionPairProto(info);
		for (const sensor of sensors) {
			const observationProto = GrpcExtensions.getObservationProto(sensor, this._observationWriter);
			agentProto.agentInfo?.observations?.push(observationProto);
		}
		const writer = protobuf.Writer.create();
		pb.communicator_objects.AgentInfoActionPairProto.encode(agentProto, writer);
		const binaryArray = writer.finish();
		this.write(binaryArray);
	}

	public close(): void {
		if (this._writerName == null) {
			return;
		}

		this.endEpisode();
		this._metaData.meanReward = this._cumulativeReward / this._metaData.numberEpisodes;
		this.writeMetadata();
		this.closeWriter();
		this._writerName = null;
	}

	private endEpisode(): void {
		this._metaData.numberEpisodes += 1;
	}
}
